"""
STL模型
test: mae; 1708.587890625， mape: 4.750349998474121, rmse: 1843.4534912109375
"""
import os
import pandas as pd
from sklearn.linear_model import LinearRegression
from statsmodels.tsa.api import STL
from statsmodels.formula.api import ols
import matplotlib.pyplot as plt
from pylab import mpl
from keras.metrics import (RootMeanSquaredError,
                           MeanAbsoluteError,
                           MeanAbsolutePercentageError)
from model_evalute import Evalute


# 指定默认字体
mpl.rcParams['font.sans-serif'] = ['SimHei']
# 解决保存图像是负号'-'显示为方块的问题
mpl.rcParams['axes.unicode_minus'] = False


def line_regression(train_data, test_data):
    """
    线性回归
    :param train_data:
    :param test_data:
    :return:
    """
    train_data_length = len(train_data)
    trend_line = LinearRegression()
    trend_line.fit(X=[[i] for i in range(train_data_length)], y=train_data)
    print(trend_line.coef_)
    print(trend_line.intercept_)

    train_trend_pred = \
        trend_line.predict(X=[[i] for i in range(train_data_length)])
    test_trend_pred = trend_line.predict(X=[[train_data_length + i]
                                            for i in range(len(test_data))])
    return train_trend_pred, test_trend_pred


def unline_regression(train_data, test_data):
    """
    非线性回归
    :param train_data:
    :param test_data:
    :return:
    """
    df = pd.DataFrame(train_data)
    df['x'] = pd.DataFrame([i for i in range(len(df))], index=df.index)
    model = ols(formula='trend ~ x + x^2', data=df).fit()
    print(model.summary())
    # pred = model.forecast(steps=len(test_data))
    train_pred = model.predict(df['x'])
    test_x = pd.Series([len(train_data) + i for i in range(len(test_data))],
                       index=test_data.index)
    test_pred = model.predict(exog=dict(x=test_x))

    return train_pred, test_pred


def stl_model(data, pic_save_path):
    """
    STL(Seasonal and Trend decomposition using Loess)分解法
    :param data:
    :param pic_save_path: 图片存储路径
    :return:
    """
    series_train = data.loc[: '2019-08']
    series_test = data.loc['2019-09':]

    model = STL(endog=series_train, trend=13, seasonal=13, robust=True).fit()

    # 查看模型
    # model.plot()
    # plt.show()

    #  强度计算
    # 残差项方差
    resid_var = model.resid.var()
    #  趋势强度
    trend_var = (model.resid + model.trend).var()
    f_trend = 1 - resid_var / trend_var
    f_trend = 0 if f_trend < 0 else f_trend

    #  季节强度
    seasonal_var = (model.resid + model.seasonal).var()
    f_seasonal = 1 - resid_var / seasonal_var
    f_seasonal = 0 if f_seasonal < 0 else f_seasonal

    print('趋势强度：{}, 季节强度：{}'.format(f_trend, f_seasonal))

    # 季节
    # 每月均值
    df_seasonal = pd.DataFrame(model.seasonal)
    df_seasonal = df_seasonal.reset_index()
    df_seasonal['month_no'] = \
        df_seasonal['month'].apply(lambda x: x.strftime('%m'))
    seasonal_mean = df_seasonal.groupby('month_no').mean()

    # 趋势
    # print(model.trend)
    # model.trend.plot()
    # plt.show()
    # 线性回归
    # train_trend_pred, test_trend_pred = \
    #     line_regression(train_data=model.trend, test_data=series_test)
    # 非线性回归
    train_trend_pred, test_trend_pred = \
        unline_regression(train_data=model.trend, test_data=series_test)

    train_trend_pred_series = \
        pd.Series(train_trend_pred, index=series_train.index)
    test_trend_pred_series = pd.Series(test_trend_pred, index=series_test.index)

    # model.seasonal.plot()
    # model.trend.plot()
    # train_trend_pred_series.plot()
    # plt.legend(['训练集',  '拟合'])
    # plt.show()

    # train预测
    train_pred = []
    for day in train_trend_pred_series.index:
        month_no = day.strftime('%m')
        train_pred.append(train_trend_pred_series.loc[day] +
                          seasonal_mean.loc[month_no]['season'])
    train_pred_series = pd.Series(train_pred, index=series_train.index)

    # test预测
    test_pred = []
    for day in test_trend_pred_series.index:
        month_no = day.strftime('%m')
        test_pred.append(test_trend_pred_series.loc[day] +
                         seasonal_mean.loc[month_no]['season'])
    test_pred_series = pd.Series(test_pred, index=series_test.index)

    #  评估
    #  整体
    data.plot()
    train_pred_series.plot()
    test_pred_series.plot()
    plt.legend(['观测值', '拟合结果', '预测结果'])
    plt.xlabel('日期')
    plt.ylabel('货运量')
    plt.title('STL模型')
    pic_name = r'STL模型.png'
    # plt.savefig(os.path.join(pic_save_path, pic_name))
    plt.show()

    # 训练
    evalute = Evalute(y_true=series_train, y_pred=train_pred_series)
    mae, mse, aic, aicc, bic = evalute.evalute_error_index()
    print('\n训练结果的参数评估：\n残差评估：mae: {}; mse: {}\n'
          '信息准则评估：aic: {}; aicc: {}; bic: {}'.
          format(mae, mse, aic, aicc, bic))

    # 测试
    evalute = Evalute(y_true=series_test, y_pred=test_pred_series)
    mae, mse, aic, aicc, bic = evalute.evalute_error_index()
    print('\n测试结果的参数评估：\n残差评估：mae: {}; mse: {}'.format(mae, mse))

    # 预测评估
    # mae
    mae = MeanAbsoluteError()
    mae.update_state(y_true=series_test, y_pred=test_pred_series)
    test_mae = mae.result()

    # mape
    mape = MeanAbsolutePercentageError()
    mape.update_state(y_true=series_test, y_pred=test_pred_series)
    test_mape = mape.result()

    # rmse
    rmse = RootMeanSquaredError()
    rmse.update_state(y_true=series_test, y_pred=test_pred_series)
    test_rmse = rmse.result()

    print('test: mae; {}， mape: {}, rmse: {}'.
          format(test_mae, test_mape, test_rmse))

    # 残差/误差图
    (train_pred_series - series_train).plot()
    (test_pred_series - series_test).plot()
    plt.legend(['残差', '误差'])
    plt.xlabel('日期')
    plt.ylabel('残差/误差')
    plt.title('残差/误差图')
    plt.show()
